from random import randint, getrandbits
from os.path import getsize

from tqdm import tqdm

def mod_pow(x, n, m):
    b = 1
    if n == 0:
        return b
    A = x
    if n & 1:
        b = x
    for i in range(1, len(bin(n))-2):
        A = (A * A) % m
        if n & (1 << i):
            b = (A * b) % m
    return b

def fermat_test(n, t):
    for i in range(t):
        a = randint(2, n - 1)
        r = mod_pow(a, n - 1, n)
        if r != 1:
            return False
    return True


def extended_gcd(a, b):
    if b == 0:
        return (a, 1, 0)
    else:
        d, x, y = extended_gcd(b, a % b)
        return (d, y, x - y * (a // b))
    

def generate_prime_candidate(length):
    p = getrandbits(length)
    p |= (1 << length - 1) | 1
    return p

def generate_prime_number(length):
    p = generate_prime_candidate(length)
    while not fermat_test(p, 128):
        p = generate_prime_candidate(length)
    return p


def generate_keypair(length):
    p = generate_prime_number(length)
    q = generate_prime_number(length)
    if abs(p - q) < 2 ** (length / 2):
        return generate_keypair(length)
    n = p * q
    phi = (p - 1) * (q - 1)
    e = 65537
    d = extended_gcd(e, phi)[1] % phi
    return ((e, n), (d, n))


def save_key(filename, keypair):
    with open(filename, 'w') as f:
        for key in keypair: 
            f.write(key.to_bytes((len(bin(key)) - 2) // 8 + 1, 'big').hex() + '\n')

def load_key(filename):
    with open(filename, 'r') as f:
        keypair = []
        for line in f:
            keypair.append(int(line, 16))
        return tuple(keypair)

def encrypt(filename, keypair, filename_out, block_size=128):
    bytes_size = getsize(filename)
    pbar = tqdm(desc=filename, total=bytes_size,  unit='B', unit_scale=True, unit_divisor=1024)
    with open(filename_out, "a") as f_e:
        f_e.write(str(bytes_size) + '\n')
    with open(filename, "rb") as f:
        data = f.read(block_size)
        while data:
            data = int.from_bytes(data, 'big')
            encrypted = mod_pow(data, keypair[0], keypair[1])
        
            with open(filename_out, "a") as f_e:
                f_e.write(encrypted.to_bytes((len(bin(encrypted)) - 2) // 8 + 1, 'big').hex() + '\n')
            data = f.read(block_size)
            pbar.update(block_size)
   

              
def decrypt(filename, keypair, filename_out, block_size=128):
    size = 0
    with open(filename, "r", encoding="ascii") as f:
        for line in f:
            if size == 0:
                size = int(line)
                pbar = tqdm(desc=filename, total=size,  unit='B', unit_scale=True, unit_divisor=1024)
                continue
            encrypted = int(line, 16)
            decrypted = mod_pow(encrypted, keypair[0], keypair[1])
            with open(filename_out, "ab") as f_d:
                if size < block_size:
                    f_d.write(decrypted.to_bytes(size, 'big'))
                else:
                    f_d.write(decrypted.to_bytes(block_size, 'big'))
                    size -= block_size
                pbar.update(block_size)